#include <io\pngstream.h>

bool 
pngstream::initializeInput(FILE *input, bool input_to_grayscale) {

  if (!validate(input)) {
    #ifdef DEBUG_INFO
	  std::cerr << "ERROR: Input file is not valid PNG file" << '\0' << std::endl;
    #endif
	  return false;
  }

  png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (png_ptr == NULL) {
    fclose(input);
    return false;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (info_ptr == NULL) {
    fclose(input);
    png_destroy_read_struct(&png_ptr, NULL, NULL);
    return false;
  }

  if (setjmp(png_jmpbuf(png_ptr))) {
    png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
    fclose(input);
    return false;
  }

  png_init_io(png_ptr, input);
  png_set_sig_bytes(png_ptr, PNGSIGSIZE);
	png_read_info(png_ptr, info_ptr);

  png_get_IHDR(png_ptr, info_ptr, &pngheader.width, &pngheader.height, &pngheader.bitdepth, &pngheader.color_type,
       &pngheader.interlace_type, NULL, NULL);

  pngheader.channels = png_get_channels(png_ptr, info_ptr);

  switch (pngheader.color_type) {
    case PNG_COLOR_TYPE_PALETTE:
      // convert palette to rgb
      png_set_palette_to_rgb(png_ptr);
      pngheader.color_type = PNG_COLOR_TYPE_RGB;

      if (input_to_grayscale) {
        // convert rgb to gray
        png_set_rgb_to_gray(png_ptr, 1, 0.0, 0.0);
        pngheader.color_type = PNG_COLOR_TYPE_GRAY;
        pngheader.channels = 1;
      }
      else pngheader.channels = 3;
    break;
  }

  return true;
}

unsigned char *
pngstream::readLine() {
  const unsigned int stride = (pngheader.width * pngheader.bitdepth * pngheader.channels) / 8;
  png_bytep rowPtr = (png_bytep) new unsigned char[stride];

  png_read_rows(png_ptr, &rowPtr, NULL, 1);
  row_index++;

  return (unsigned char *) rowPtr;
}

bool 
pngstream::hasNext() {
  return (row_index < pngheader.height);
};

void 
pngstream::finishInput() {
  png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
};

bool 
pngstream::validate(FILE *fp) {
  //Allocate a buffer of 8 bytes, where we can put the file signature.
  char signature[PNGSIGSIZE];

  if (fp) {
    if (fread(signature, 1, PNGSIGSIZE, fp) != PNGSIGSIZE) return 0;

    return(!png_sig_cmp((png_bytep) signature, (png_size_t)0, PNGSIGSIZE));
  }
	return false;
}

bool 
pngstream::initializeOutput(FILE *output, int width, int height, int bit_depth, int color_type) {

  png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (png_ptr == NULL) {
    fclose(output);
    return false;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (info_ptr == NULL) {
    png_destroy_write_struct(&png_ptr, NULL);
    return false;
  }

  if (setjmp(png_jmpbuf(png_ptr))) {
    png_destroy_write_struct(&png_ptr, &info_ptr);
    return false;
  }

  png_init_io(png_ptr, output);

  if (setjmp(png_jmpbuf(png_ptr))) {
    png_destroy_write_struct(&png_ptr, &info_ptr);
    return false;
  }

  png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, color_type,
               PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

  png_write_info(png_ptr, info_ptr);

  pngheader.channels = png_get_channels(png_ptr, info_ptr);

  return true;
}

bool 
pngstream::write(unsigned char **data) {

  if (setjmp(png_jmpbuf(png_ptr))) {
		png_destroy_write_struct(&png_ptr, &info_ptr);
    return false;
  }

	png_write_image(png_ptr, data);

  if (setjmp(png_jmpbuf(png_ptr))) {
		png_destroy_write_struct(&png_ptr, &info_ptr);
    return false;
  }

	png_write_end(png_ptr, NULL);

  return true;
}

void 
pngstream::finishOutput() {
  png_destroy_write_struct(&png_ptr, &info_ptr);
};
